import os
from typing import List

from openai import OpenAI
from contextlib import contextmanager
from qdrant_client import QdrantClient
from qdrant_client.http import models

# 从环境变量中获取阿里云-百练的API Key
DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY")
# 阿里云-百练的官网地址
DASHSCOPE_API_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"

qwen_client = OpenAI(api_key=DASHSCOPE_API_KEY, base_url=DASHSCOPE_API_BASE_URL)


def get_text_embedding(text: str, model="text-embedding-v1"):
    """
    获取文本的嵌入向量
    :param text:
    :param model:
    :return:
    """
    embedding_response = qwen_client.embeddings.create(input=text, model=model)
    return embedding_response.data[0].embedding


def get_texts_embedding(texts: List[str], model="text-embedding-v1"):
    """
    获取文本的嵌入向量
    :param texts:
    :param model:
    :return:
    """
    embedding_response = qwen_client.embeddings.create(input=texts, model=model)
    return [item.embedding for item in embedding_response.data]


@contextmanager
def get_qdrant_client(host="localhost", port=6333):
    qdrant_client = QdrantClient(host=host, port=port)
    try:
        yield qdrant_client
    finally:
        qdrant_client.close()


def create_collection(qdrant_client, collection_name="text_collection", vector_size=1536):
    """
    创建文本向量集合
    :param qdrant_client:
    :param collection_name:
    :param vector_size:
    :return:
    """

    result = qdrant_client.create_collection(
        collection_name=collection_name,
        vectors_config=models.VectorParams(
            size=vector_size,  # 向量维度 根据使用的模型确定
            distance=models.Distance.COSINE  # 距离度量方式: 余弦相似度
        )
    )

    return result


def insert_single_point(qdrant_client, collection_name, point_id, vector, payload):
    """
    插入单个点数据
    :param qdrant_client:
    :param collection_name:
    :param point_id:
    :param vector:
    :param payload:
    :return:
    """
    point = models.PointStruct(
        id=point_id,
        vector=vector,
        payload=payload
    )
    result = qdrant_client.upsert(
        collection_name=collection_name,
        points=[point]
    )
    return result


def search_similar_vectors(qdrant_client, collection_name, query_vector, limit=1):
    """
    搜索相似向量
    :param qdrant_client:
    :param collection_name:
    :param query_vector:
    :param limit:
    :return:
    """
    results = qdrant_client.search(
        collection_name=collection_name,
        query_vector=query_vector,
        limit=limit
    )
    return results


def search_with_filter(client, collection_name, query_vector, filter_conditions: list[dict], limit=1):
    """
    带过滤条件的搜索
    :param client:
    :param collection_name:
    :param query_vector:
    :param filter_conditions:
    :param limit:
    :return:
    """
    filter = models.Filter(
        must=[
            models.FieldCondition(
                key=condition["key"],
                match=models.MatchValue(value=condition["value"])
            )
            for condition in filter_conditions
        ]
    )
    results = client.search(
        collection_name=collection_name,
        query_vector=query_vector,
        query_filter=filter,
        limit=limit
    )
    return results


if __name__ == '__main__':
    text = ["AI学习"]
    embedding = get_text_embedding(text)
    if embedding:
        vector_size = len(embedding)
        print(f"嵌入模型输出向量维度: {vector_size}")
        print(embedding[:10])

    with get_qdrant_client() as qdrant_client:
        collection_name = "ai_collection"
        # 判断集合是否存在
        is_exist = qdrant_client.collection_exists(collection_name=collection_name)
        if not is_exist:
            # 创建集合
            create_collection(qdrant_client, collection_name, vector_size)
        else:
            # 删除集合
            qdrant_client.delete_collection(collection_name=collection_name)
            create_collection(qdrant_client, collection_name, vector_size)

        # 插入文本数据
        embedding1 = get_text_embedding(["机器学习&深度学习"])
        result1 = insert_single_point(qdrant_client, collection_name, point_id=1, vector=embedding1,
                                      payload={"text": "机器学习&深度学习", "category": "AI", "date": "2025-06-06"})
        print(f"result1 ==> {result1}")

        # 插入带标签的数据
        embedding2 = get_text_embedding(["深度学习是AI的核心技术"])
        result2 = insert_single_point(qdrant_client, collection_name, point_id=2, vector=embedding2,
                                      payload={"text": "深度学习是AI的核心技术", "category": "科技", "tags": ["AI", "ML", "深度学习"],
                                               "author": "yudao"})
        print(f"result2 ==> {result2}")

        # 插入带评分的数据
        embedding3 = get_text_embedding(["AI的最新进展"])
        result3 = insert_single_point(qdrant_client, collection_name, point_id=3, vector=embedding3,
                                      payload={"text": "AI的最新进展", "category": "AI", "score": 0.95, "source": "学术论文"})
        print(f"result2 ==> {result3}")

        # 搜索相似向量
        embedding4 = get_text_embedding(["AI人工智能技术的最新进展"])
        result4 = search_similar_vectors(qdrant_client, collection_name, query_vector=embedding4, limit=2)
        print(f"result4 ==> {result4}")

        # 按类别过滤
        filter_conditions = [{"key": "category", "value": "AI"}]
        result5 = search_with_filter(qdrant_client, collection_name, query_vector=get_text_embedding(["AI"]),
                                     filter_conditions=filter_conditions, limit=2)
        print(f"result5 ==> {result5}")

        # 按类别和日期过滤
        filter_conditions = [{"key": "category", "value": "AI"}, {"key": "date", "value": "2025-06-06"}]
        result6 = search_with_filter(qdrant_client, collection_name, query_vector=get_text_embedding(["AI"]),
                                     filter_conditions=filter_conditions, limit=2)
        print(f"result6 ==> {result6}")
